{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Text to Multiclass Explanation: Language Modeling Example"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This notebook demostrates how to get explanations for the top-k next words generated by a language model. In this demo, we use the pretrained gpt2 model provided by hugging face (https://huggingface.co/gpt2) to predict the top-k next words. By looking at the top-k next words, we treat them as k separate classes and then learn the explanations for each of these k words. We thereby are able to explain the contribution of words in the input that are responsible for the likelihood of the top-k next words to be predicted. "
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"\n",
"import shap"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Load model and tokenizer"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"tokenizer = AutoTokenizer.from_pretrained(\"gpt2\", use_fast=True)\n",
"model = AutoModelForCausalLM.from_pretrained(\"gpt2\").cuda()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We next wrap the model with the TopKLM model which extracts the log odds of the top-k next words and also create a Text masker by initializing it with the mask_token = \"...\" and set collapse_mask_token = True, which is used for infilling text during perturbation of the inputs."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"wrapped_model = shap.models.TopKLM(model, tokenizer, k=100)\n",
"masker = shap.maskers.Text(tokenizer, mask_token=\"...\", collapse_mask_token=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define data"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here we set the initial text for which we want the gpt2 model to predict the next word"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"s = [\"In a shocking finding, scientists discovered a herd of unicorns living in a\"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Create explainer object"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"explainer = shap.Explainer(wrapped_model, masker)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Compute SHAP values"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"shap_values = explainer(s)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Visualize the SHAP values across the input sentence for the top-k next words"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can now see the top-k next words predicted by gpt2 under \"Output Text\" in the plot below and hover over each of the token to understand which words in the input sentence are driving the generation of the particular output word to be predicted"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
" \n",
"\n",
"